"""
工具函数
"""
import cv2
import numpy as np
import os
from PIL import Image, ImageDraw, ImageFont

def draw_contours(image, contours, shape_name):
    """在图像上绘制轮廓和标签"""
    # 复制图像以避免修改原图
    result_image = image.copy()
    
    # 绘制轮廓
    cv2.drawContours(result_image, contours, -1, (0, 255, 0), 2)
    
    # 为每个轮廓添加标签
    for i, contour in enumerate(contours):
        # 计算轮廓中心
        M = cv2.moments(contour)
        if M["m00"] != 0:
            cx = int(M["m10"] / M["m00"])
            cy = int(M["m01"] / M["m00"])
            
            # 添加文本标签
            label = f"{shape_name}_{i+1}"
            cv2.putText(result_image, label, (cx-20, cy-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
    
    return result_image

def calculate_shape_properties(contour):
    """计算图形属性"""
    # 面积
    area = cv2.contourArea(contour)
    
    # 周长
    perimeter = cv2.arcLength(contour, True)
    
    # 边界矩形
    x, y, w, h = cv2.boundingRect(contour)
    
    # 长宽比
    aspect_ratio = float(w) / h if h != 0 else 0
    
    # 矩形度（轮廓面积与边界矩形面积的比值）
    rectangularity = float(area) / (w * h) if w * h != 0 else 0
    
    # 圆形度
    circularity = 4 * np.pi * area / (perimeter * perimeter) if perimeter != 0 else 0
    
    # 凸包
    hull = cv2.convexHull(contour)
    hull_area = cv2.contourArea(hull)
    
    # 凸性（轮廓面积与凸包面积的比值）
    convexity = float(area) / hull_area if hull_area != 0 else 0
    
    return {
        "area": area,
        "perimeter": perimeter,
        "aspect_ratio": aspect_ratio,
        "rectangularity": rectangularity,
        "circularity": circularity,
        "convexity": convexity,
        "bounding_rect": (x, y, w, h)
    }

def create_demo_images():
    """创建演示用的图形图像"""
    demo_dir = "assets/demo_images"
    os.makedirs(demo_dir, exist_ok=True)
    
    # 图像尺寸
    width, height = 300, 300
    
    # 创建三角形
    triangle_img = Image.new('RGB', (width, height), 'white')
    draw = ImageDraw.Draw(triangle_img)
    triangle_points = [(150, 50), (75, 200), (225, 200)]
    draw.polygon(triangle_points, outline='black', fill='lightblue', width=3)
    triangle_img.save(os.path.join(demo_dir, 'triangle.png'))
    
    # 创建正方形
    square_img = Image.new('RGB', (width, height), 'white')
    draw = ImageDraw.Draw(square_img)
    draw.rectangle([75, 75, 225, 225], outline='black', fill='lightgreen', width=3)
    square_img.save(os.path.join(demo_dir, 'square.png'))
    
    # 创建长方形
    rectangle_img = Image.new('RGB', (width, height), 'white')
    draw = ImageDraw.Draw(rectangle_img)
    draw.rectangle([50, 100, 250, 200], outline='black', fill='lightcoral', width=3)
    rectangle_img.save(os.path.join(demo_dir, 'rectangle.png'))
    
    # 创建圆形
    circle_img = Image.new('RGB', (width, height), 'white')
    draw = ImageDraw.Draw(circle_img)
    draw.ellipse([75, 75, 225, 225], outline='black', fill='lightyellow', width=3)
    circle_img.save(os.path.join(demo_dir, 'circle.png'))
    
    # 创建混合图形
    mixed_img = Image.new('RGB', (width, height), 'white')
    draw = ImageDraw.Draw(mixed_img)
    # 小三角形
    draw.polygon([(75, 75), (50, 125), (100, 125)], outline='red', fill='pink', width=2)
    # 小正方形
    draw.rectangle([150, 50, 200, 100], outline='blue', fill='lightblue', width=2)
    # 小圆形
    draw.ellipse([200, 150, 250, 200], outline='green', fill='lightgreen', width=2)
    mixed_img.save(os.path.join(demo_dir, 'mixed_shapes.png'))
    
    print(f"演示图像已创建在 {demo_dir} 目录")

def resize_image(image, max_size=800):
    """调整图像大小，保持长宽比"""
    height, width = image.shape[:2]
    
    if max(height, width) <= max_size:
        return image
    
    if height > width:
        new_height = max_size
        new_width = int(width * max_size / height)
    else:
        new_width = max_size
        new_height = int(height * max_size / width)
    
    resized = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)
    return resized

def enhance_image_contrast(image, alpha=1.5, beta=0):
    """增强图像对比度"""
    enhanced = cv2.convertScaleAbs(image, alpha=alpha, beta=beta)
    return enhanced

def remove_noise(image, kernel_size=5):
    """去除图像噪声"""
    # 使用形态学操作去噪
    kernel = np.ones((kernel_size, kernel_size), np.uint8)
    
    # 开运算（先腐蚀后膨胀）去除小的噪声点
    opening = cv2.morphologyEx(image, cv2.MORPH_OPEN, kernel)
    
    # 闭运算（先膨胀后腐蚀）填充图形内部的小洞
    closing = cv2.morphologyEx(opening, cv2.MORPH_CLOSE, kernel)
    
    return closing

def get_shape_color(shape_name):
    """根据图形名称获取对应的颜色"""
    color_map = {
        "triangle": (255, 0, 0),    # 红色
        "square": (0, 255, 0),      # 绿色
        "rectangle": (0, 0, 255),   # 蓝色
        "circle": (255, 255, 0),    # 黄色
        "unknown": (128, 128, 128)  # 灰色
    }
    return color_map.get(shape_name.lower(), (128, 128, 128))

def draw_detection_info(image, shapes):
    """在图像上绘制检测信息统计"""
    if not shapes:
        return image
    
    result_image = image.copy()
    
    # 统计各种图形的数量
    shape_counts = {}
    for shape in shapes:
        name = shape["name"]
        shape_counts[name] = shape_counts.get(name, 0) + 1
    
    # 在图像顶部绘制统计信息
    y_offset = 30
    for shape_name, count in shape_counts.items():
        text = f"{shape_name}: {count}"
        color = get_shape_color(shape_name)
        
        cv2.putText(result_image, text, (10, y_offset), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)
        y_offset += 30
    
    return result_image

def validate_image_format(image_path):
    """验证图像格式是否支持"""
    supported_formats = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']
    
    if not os.path.exists(image_path):
        return False, "文件不存在"
    
    file_ext = os.path.splitext(image_path)[1].lower()
    if file_ext not in supported_formats:
        return False, f"不支持的文件格式: {file_ext}"
    
    try:
        # 尝试读取图像
        img = cv2.imread(image_path)
        if img is None:
            return False, "无法读取图像文件"
    except Exception as e:
        return False, f"读取图像时出错: {str(e)}"
    
    return True, "图像格式有效"

def save_detection_result(image, shapes, output_path):
    """保存检测结果到文件"""
    try:
        # 在图像上绘制检测结果
        result_image = draw_detection_info(image, shapes)
        
        # 保存图像
        cv2.imwrite(output_path, result_image)
        
        # 保存检测数据到文本文件
        txt_path = os.path.splitext(output_path)[0] + '_results.txt'
        with open(txt_path, 'w', encoding='utf-8') as f:
            f.write(f"检测结果 - 总共发现 {len(shapes)} 个图形:\n\n")
            for i, shape in enumerate(shapes, 1):
                f.write(f"{i}. {shape['name']}\n")
                f.write(f"   置信度: {shape['confidence']:.3f}\n")
                f.write(f"   面积: {shape['area']:.2f}\n")
                f.write(f"   周长: {shape['perimeter']:.2f}\n")
                f.write(f"   长宽比: {shape['aspect_ratio']:.3f}\n")
                f.write(f"   圆形度: {shape['circularity']:.3f}\n")
                f.write("\n")
        
        return True, f"结果已保存到 {output_path} 和 {txt_path}"
    
    except Exception as e:
        return False, f"保存失败: {str(e)}"

def create_shape_template(shape_name, size=(100, 100)):
    """创建标准图形模板"""
    width, height = size
    template = np.zeros((height, width, 3), dtype=np.uint8)
    template.fill(255)  # 白色背景
    
    center_x, center_y = width // 2, height // 2
    
    if shape_name.lower() == "circle":
        cv2.circle(template, (center_x, center_y), min(width, height) // 3, (0, 0, 0), 2)
    
    elif shape_name.lower() == "square":
        size = min(width, height) // 3
        top_left = (center_x - size, center_y - size)
        bottom_right = (center_x + size, center_y + size)
        cv2.rectangle(template, top_left, bottom_right, (0, 0, 0), 2)
    
    elif shape_name.lower() == "triangle":
        points = np.array([
            [center_x, center_y - height // 3],
            [center_x - width // 3, center_y + height // 4],
            [center_x + width // 3, center_y + height // 4]
        ], np.int32)
        cv2.polylines(template, [points], True, (0, 0, 0), 2)
    
    elif shape_name.lower() == "rectangle":
        width_offset = width // 3
        height_offset = height // 4
        top_left = (center_x - width_offset, center_y - height_offset)
        bottom_right = (center_x + width_offset, center_y + height_offset)
        cv2.rectangle(template, top_left, bottom_right, (0, 0, 0), 2)
    
    return template

def calculate_detection_accuracy(true_shapes, detected_shapes):
    """计算检测准确率"""
    if not true_shapes:
        return 1.0 if not detected_shapes else 0.0
    
    if not detected_shapes:
        return 0.0
    
    # 简单的名称匹配计算准确率
    true_names = [shape.lower() for shape in true_shapes]
    detected_names = [shape["name"].lower() for shape in detected_shapes]
    
    correct_count = 0
    for true_name in true_names:
        if true_name in detected_names:
            correct_count += 1
            detected_names.remove(true_name)  # 避免重复计算
    
    precision = correct_count / len(detected_shapes) if detected_shapes else 0
    recall = correct_count / len(true_shapes) if true_shapes else 0
    
    f1_score = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0
    
    return {
        "precision": precision,
        "recall": recall,
        "f1_score": f1_score,
        "correct_count": correct_count,
        "total_true": len(true_shapes),
        "total_detected": len(detected_shapes)
    }